#include "Collision1D_Coulomb.h"
#include "../Parallel/ParallelMPI.h"
#include "../Field/Field1D.h"
#include "../Field/Field2D.h"
#include "../PhysicalField/PhysicalField.h"
#include "../Species/Species1D.h"


#include <cmath>
#include <iomanip>
#include <algorithm>
#include <ostream>
#include <sstream>

using namespace std;


Collision1D_Coulomb::Collision1D_Coulomb(PicParams* params_, CollisionParameter collision_param_, int collision_number_):
    Collision1D(params_, collision_param_, collision_number_)
{
    twoPi = 2.0 * params->const_pi;
    e_ov_ephi0 = params->const_e / params->const_ephi0;

    n_species = collision_param.species_number_list.size();
    n_particle.resize(n_species);
    density.resize(n_species);
    index.resize(n_species);
    n_particle_begin.resize(n_species);

    debye_length.resize(params->dims_space[0]);

    cal_collision_pairs();
}

Collision1D_Coulomb::~Collision1D_Coulomb()
{

}


void Collision1D_Coulomb::collide(PhysicalField* fields, vector<Species*>& vecSpecies, Diagnostic* diag, int itime)
{
    double gx, gy, gz, g_magnitude, g_square, g_3, g_p, hx, hy, hz, s, A12, cosX, sinX, phi;
    int i1, i2;
    double time_zoom_factor;

    //cout<<"coulomb collision 00000"<<endl;
    if(collision_param.timestep_collision == 0 || itime % collision_param.timestep_collision != 0)
    {
        return;
    }

    cal_log_coulomb(fields);

    //cout<<"coulomb collision 11111"<<endl;
    for(int i_space = 0 ; i_space < params->dims_space[0] ; i_space++)
    {
        time_zoom_factor = collision_param.time_zoom_factor;
        //determine if zoom collision frequency
        if(collision_param.region_collision_zoom_start.size() > 0 && collision_param.region_collision_zoom_end.size() > 0)
        {
            if(params->local_min[0] + (i_space + 1) * params->cell_length[0] < collision_param.region_collision_zoom_start[0]
               || params->local_min[0] + i_space * params->cell_length[0] > collision_param.region_collision_zoom_end[0])
            {
                time_zoom_factor = 1.0;
            }
        }

        density_all = 0;
        for(int i_species = 0; i_species < n_species; i_species++)
        {
          Species1D* species1D = static_cast<Species1D*>(vecSpecies[ collision_param.species_number_list[i_species] ]);
          double weight = species1D->species_param.weight;
          n_particle[i_species] = species1D->particles_[i_space]->size();
          density[i_species] = n_particle[i_species] * weight;
          density_all += density[i_species];
          n_particle_begin[i_species] = 0;
        }
        if(density_all == 0.0)
        {
          continue;
        }
        density_all_inv = 1.0 / density_all;

        for(int i_species = 0; i_species < n_species; i_species++)
        {
            index[i_species].resize(n_particle[i_species]);
            for (int j = 0; j < index[i_species].size(); j++)
            {
              index[i_species][j] = j;
            }
            random_shuffle(index[i_species].begin(), index[i_species].end());
         }

        //cout<<"coulomb collision 2222 "<<i_space<<endl;
        for(int i_pair = 0; i_pair < n_collision_pairs - n_species; i_pair++)
        {
            int is1 = iSpec1[i_pair];
            int is2 = iSpec2[i_pair];
            //!!!it may be wrong here, is1 correspond to Species or Collision species?
            n_indexes[i_pair] = n_particle[is1] * density[is2] * density_all_inv;
            isOdd[i_pair] = false;

            n_particle_begin[is1] += n_indexes[i_pair];
            n_particle_begin[is2] += n_indexes[i_pair];

        }

        for(int i_pair = n_collision_pairs - n_species; i_pair < n_collision_pairs; i_pair++)
        {
            int is1 = iSpec1[i_pair];

            n_indexes[i_pair] = n_particle[is1] - n_particle_begin[is1];

            if(n_indexes[i_pair] % 2 == 0)
            {
                isOdd[i_pair] = false;
                n_indexes[i_pair] /= 2;
            }
            else
            {
                isOdd[i_pair] = true;
                n_indexes[i_pair] = (n_indexes[i_pair] + 1) / 2;
                int i_species = i_pair - n_collision_pairs + n_species;
                index[i_species].push_back(n_particle[i_species] - 1);
            }
        }


        for(int i_species = 0; i_species < n_species; i_species++)
        {
          n_particle_begin[i_species] = 0;
        }

        //cout<<"coulomb collision 3333 "<<n_indexes[0]<<endl;
        //loop on collision_pairs
        for(int i_pair = 0; i_pair < n_collision_pairs; i_pair++)
        {
          //cout<<"n_collision_pairs "<<n_collision_pairs<<"  "<<i_pair<<endl;
          int is1 = iSpec1[i_pair];
          int is2 = iSpec2[i_pair];
          Species1D* s1 = static_cast<Species1D*>(vecSpecies[ is1 ]);
          Species1D* s2 = static_cast<Species1D*>(vecSpecies[ is2 ]);
          auto cell_particles1 = s1->particles_[i_space];
          auto cell_particles2 = s2->particles_[i_space];

          if(log_coulomb[i_pair][i_space] == 0.0)
          {
            continue;
          }
          A12 = gamma2[i_pair] * density_all * log_coulomb[i_pair][i_space];


          //loop on collision pairs of particles
          for(int i_index = 0; i_index < n_indexes[i_pair]; i_index++)
          {
              //cout<<"n_indexes[i] "<<n_indexes[i_pair]<<"  "<<i_index<<endl;
              if(isIntra[i_pair])
              {
                  //if(i_index + n_particle_begin[is1] > n_particle[is1]) {cout<<"is1 error "<<endl;}
                  //if(i_index + n_particle_begin[is2] + n_indexes[i_pair] > n_particle[is2]) {cout<<"is2 error "<<endl;}
                  i1 = index[is1][i_index + n_particle_begin[is1]];
                  i2 = index[is2][i_index + n_particle_begin[is2] + n_indexes[i_pair]];
              }
              else
              {
                i1 = index[is1][i_index + n_particle_begin[is1]];
                i2 = index[is2][i_index + n_particle_begin[is2]];
              }

              //if(i1 >= cell_particles1->size()) { cout<<"i1 error "<<n_particle[is1]<<"  "<<i1<<" "<<n_indexes[0]<<"  "<<n_indexes[1]<<"  "<<n_indexes[2]<<endl; }
              //if(i2 >= cell_particles2->size()) { cout<<"i2 error "<<n_particle[is2]<<" "<<i2<<"  "<<n_indexes[0]<<"  "<<n_indexes[1]<<"  "<<n_indexes[2]<<endl;  }
              auto& p1 = cell_particles1->data_[i1];
              auto& p2 = cell_particles2->data_[i2];

              //Calculate stuff
              gx = p1.velocity[0] - p2.velocity[0];
              gy = p1.velocity[1] - p2.velocity[1];
              gz = p1.velocity[2] - p2.velocity[2];
              g_square = gx*gx + gy*gy + gz*gz;
              g_magnitude = sqrt( g_square );
              g_3 = g_square * g_magnitude;
              g_p = sqrt( gy*gy + gz*gz );
              if(g_p == 0.0 || g_3 == 0.0)
              {
                  continue;
              }
              //the formula below the equation (101)
              s = A12 * collision_param.timestep_collision * params->dt * time_zoom_factor / g_3;
              if(isOdd[i_pair])
              {
                int n_temp = n_indexes[i_pair] * 2 -1;
                s *= ( n_temp / (n_temp+1) );
              }

              //Pick the deflection angles according to Nanbu's theory
              //ref: improved modeing of relativistic collisions and collisional ionization in paritcle in cell codes
              cosX = cos_chi(s);
              if(cosX < -1.0)
              {
                cosX = -1.0;
                cout<<"cosX < -1.0"<<endl;
              }
              else if(cosX > 1.0)
              {
                cosX = 1.0;
                cout<<"cosX > 1.0"<<endl;
              }

              if(cosX*cosX > 1.0)
              {
                cout<<"cosX * cosX > 1.0"<<endl;
                continue;
              }
              sinX = sqrt( 1. - cosX*cosX );


              //todo make a faster rand by preallocating ??
              phi = twoPi * ((double)rand() / RAND_MAX);
              hx = g_p * cos(phi);
              hy = -( gx * gy * cos(phi) + g_magnitude * gz * sin(phi) ) / g_p;
              hz = -( gx * gz * cos(phi) - g_magnitude * gy * sin(phi) ) / g_p;

              //double total_energy_before = 0.5 * s1->species_param.mass * (p1.velocity[0] * p1.velocity[0] + p1.velocity[1] * p1.velocity[1] + p1.velocity[2] * p1.velocity[2])
              //                          + 0.5 * s2->species_param.mass * (p2.velocity[0] * p2.velocity[0] + p2.velocity[1] * p2.velocity[1] + p2.velocity[2] * p2.velocity[2]);

              p1.velocity[0] = p1.velocity[0] - mr2[i_pair] * ( gx * (1-cosX) + hx * sinX );
              p1.velocity[1] = p1.velocity[1] - mr2[i_pair] * ( gy * (1-cosX) + hy * sinX );
              p1.velocity[2] = p1.velocity[2] - mr2[i_pair] * ( gz * (1-cosX) + hz * sinX );

              p2.velocity[0] = p2.velocity[0] + mr1[i_pair] * ( gx * (1-cosX) + hx * sinX );
              p2.velocity[1] = p2.velocity[1] + mr1[i_pair] * ( gy * (1-cosX) + hy * sinX );
              p2.velocity[2] = p2.velocity[2] + mr1[i_pair] * ( gz * (1-cosX) + hz * sinX );

              //double total_energy_after = 0.5 * s1->species_param.mass * (p1.velocity[0] * p1.velocity[0] + p1.velocity[1] * p1.velocity[1] + p1.velocity[2] * p1.velocity[2])
              //                          + 0.5 * s2->species_param.mass * (p2.velocity[0] * p2.velocity[0] + p2.velocity[1] * p2.velocity[1] + p2.velocity[2] * p2.velocity[2]);
              //cout<<"coloumb collision energy change: "<<total_energy_before<<"  "<<total_energy_after<<"  "<<total_energy_after - total_energy_before<<endl;
          }
          n_particle_begin[is1] += n_indexes[i_pair];
          n_particle_begin[is2] += n_indexes[i_pair];
        }

    }
    //cout<<"coulomb collision 66666"<<endl;
}


void Collision1D_Coulomb::cal_collision_pairs()
{
  int is1, is2;
  double m1, m2;
  double charge1, charge2;
  n_collision_pairs = n_species + n_species * (n_species - 1) / 2;

  n_indexes.resize(n_collision_pairs);
  gamma1.resize(n_collision_pairs);
  gamma2.resize(n_collision_pairs);
  mr1.resize(n_collision_pairs);
  mr2.resize(n_collision_pairs);
  m12.resize(n_collision_pairs);
  m1_inv.resize(n_collision_pairs);
  m2_inv.resize(n_collision_pairs);
  isOdd.resize(n_collision_pairs);

  log_coulomb.resize(n_collision_pairs);
  for(int i_pair = 0; i_pair < n_collision_pairs; i_pair++)
  {
      log_coulomb[i_pair].resize(params->dims_space[0]);
  }

  for(int i = 0; i < n_species - 1; i++)
  {
    for(int j = i + 1; j < n_species; j++)
    {
      iSpec1.push_back(collision_param.species_number_list[i]);
      iSpec2.push_back(collision_param.species_number_list[j]);
      isIntra.push_back(false);
    }
  }

  for(int i = 0; i < n_species; i++)
  {
    iSpec1.push_back(collision_param.species_number_list[i]);
    iSpec2.push_back(collision_param.species_number_list[i]);
    isIntra.push_back(true);
  }

  for(int i = 0; i < n_collision_pairs; i++)
  {
    is1 = iSpec1[i];
    is2 = iSpec2[i];
    m1 = params->species_param[is1].mass;
    m2 = params->species_param[is2].mass;
    charge1 = params->species_param[is1].charge;
    charge2 = params->species_param[is2].charge;
    m1_inv[i] = 1.0 / m1;
    m2_inv[i] = 1.0 / m2;

    //used in equation (104a) and (104b)
    mr1[i] = m1 / (m1 + m2);
    mr2[i] = m2 / (m1 + m2);
    //the reduced mass
    m12[i] = m1 * m2 / (m1 + m2);
    //used in equation (96)
    gamma1[i] = 4.0 * params->const_pi * params->const_ephi0 * m12[i] / abs(charge1 * charge2);
    //used in equation (95)
    gamma2[i] = 4.0 * params->const_pi / ( gamma1[i] * gamma1[i] );
  }

}


//equation 96
void Collision1D_Coulomb::cal_log_coulomb(PhysicalField* fields)
{
  double g12_square;
  int iS1, iS2;
  double T1, T2, Vx, Vy, Vz;

  //calculate electron debye_length
  iS1 = 0;
  Field1D<double>* T_electron = static_cast<Field1D<double>*>(fields->T_s[ iS1 ]);
  Field1D<double>* density_electron = static_cast<Field1D<double>*>(fields->rho_s[ iS1 ]);
  for(int i_space = 0; i_space < params->dims_space[0]; i_space++)
  {
    if( (*density_electron)( i_space + params->oversize[0] ) == 0.0 )
    {
      debye_length[i_space] = 0.0;
    }
    else
    {
      debye_length[i_space] = sqrt(params->const_ephi0 * (*T_electron)( i_space + params->oversize[0] ) / ( (*density_electron)( i_space + params->oversize[0] ) * params->const_e));
    }
  }

  for(int i = 0; i < n_collision_pairs; i++)
  {
    iS1 = iSpec1[i];
    iS2 = iSpec2[i];
    Field1D<double>* T1_s = static_cast<Field1D<double>*>(fields->T_s[ iS1 ]);
    Field1D<double>* T2_s = static_cast<Field1D<double>*>(fields->T_s[ iS2 ]);

    Field1D<double>* Vx1_s = static_cast<Field1D<double>*>(fields->Vx_s[ iS1 ]);
    Field1D<double>* Vy1_s = static_cast<Field1D<double>*>(fields->Vy_s[ iS1 ]);
    Field1D<double>* Vz1_s = static_cast<Field1D<double>*>(fields->Vz_s[ iS1 ]);

    Field1D<double>* Vx2_s = static_cast<Field1D<double>*>(fields->Vx_s[ iS2 ]);
    Field1D<double>* Vy2_s = static_cast<Field1D<double>*>(fields->Vy_s[ iS2 ]);
    Field1D<double>* Vz2_s = static_cast<Field1D<double>*>(fields->Vz_s[ iS2 ]);

    for(int i_space = 0; i_space < params->dims_space[0]; i_space++)
    {
      T1 = (*T1_s)(i_space + params->oversize[0]) * 3.0 * params->const_e * m1_inv[i];
      T2 = (*T2_s)(i_space + params->oversize[0]) * 3.0 * params->const_e * m2_inv[i];
      Vx = (*Vx1_s)(i_space + params->oversize[0]) - (*Vx2_s)(i_space + params->oversize[0]);
      Vy = (*Vy1_s)(i_space + params->oversize[0]) - (*Vy2_s)(i_space + params->oversize[0]);
      Vz = (*Vz1_s)(i_space + params->oversize[0]) - (*Vz2_s)(i_space + params->oversize[0]);
      //the formula below equation (96)
      g12_square = T1 + T2 + Vx*Vx + Vy*Vy + Vz*Vz;
      if(debye_length[i_space] == 0.0 || g12_square == 0.0)
      {
            log_coulomb[i][i_space] = 0.0;
      }
      else
      {
          log_coulomb[i][i_space] = log( gamma1[i] * g12_square * debye_length[i_space] );
          //log_info<<"log_coulomb[i][i_space]: "<<log_coulomb[i][i_space];
          //log_coulomb[i][i_space] = 17.0;
      }

    }
  }

}



//Technique given by Nanbu in http://dx.doi.org/10.1103/PhysRevE.55.4642
//  to pick randomly the deflection angle cosine, in the center-of-mass frame.
//It involves the "s" parameter (~ collision frequency * deflection expectation)
//  and a random number "U".
//Technique slightly modified in http://dx.doi.org/10.1063/1.4742167
//ref: improved modeing of relativistic collisions and collisional ionization in paritcle in cell codes
inline double Collision1D_Coulomb::cos_chi(double s)
{

    double A, invA;
    //!\todo make a faster rand by preallocating ??
    double U = (double)rand() / RAND_MAX;

    if( s < 0.1 ) {
        //if ( U<0.0001 ) U=0.0001; //ensures cos_chi > 0
        while(U == 0 || 1. + s*log(U) <= 0.0)
        {
          U = (double)rand() / RAND_MAX;
        }
        return 1. + s*log(U);
    }
    if( s < 3.  ) {
        //the polynomial has been modified from the article in order to have a better form
        invA = 0.00569578 +(0.95602 + (-0.508139 + (0.479139 + ( -0.12789 + 0.0238957*s )*s )*s )*s )*s;
        A = 1./invA;
        return  invA  * log( exp(-A) + 2.*U*sinh(A) );
    }
    if( s < 6.  ) {
        A = 3.*exp(-s);
        return (1./A) * log( exp(-A) + 2.*U*sinh(A) );
    }
    return 2.*U - 1.;

}
